#!/usr/bin/env python

''' -------------------------------------------------------------------
    This script is an experiment, to try and see if one can simply rely
    on the pseudo-inverse of the kernel-phase matrix to reconstruct a
    phase-map from a set of kernel-phases in the hope that this map can
    then be turned into an image.

    Preliminary results aren't all that great, but I include the script
    as part of the software release, as I think it is still instructive,
    and hasn't revealed all there is to know about this.

    Frantz (February 2013)

    After including the change in the code so that it now directly
    manipulates PHI and not RED*PHI like in the first version, the
    results look much more encouraging!

    The pseudo inverse approach may very well be the right approach
    to the imaging problem after all...

    Frantz (February 2013)
    ------------------------------------------------------------------- '''

import pysco
import numpy as np
import matplotlib.pyplot as plt
import pyfits as pf
import pdb

import pysco.fitting as fit
import pysco.core as pc

from scipy.interpolate import griddata

geometry = "hst"

# 0. for reference: make a kernel-phase data structure from coord file
# --------------------------------------------------------------------
a = pysco.KPI("geometry/"+geometry+".txt")
a.name = "Thick single spider arm pupil"
a.save_to_file(geometry+".kpi.gz")

# 1. load the dataset
# -------------------

ddir = "/Users/frantz/prog/python/optics copy/spk_ctrl/"

cal = pysco.KPO(geometry+".kpi.gz")
cal.extract_KPD(ddir + "undersamp.fits", ave="none", re_center=False, wfs=True)


# determine the pseudo inverse for image reconstruction
# -----------------------------------------------------
ne = 288 # number of modes reconstructed for the inverse
U, s, Vh = np.linalg.svd(cal.kpi.KerPhi, full_matrices=1)
S = np.zeros_like(cal.kpi.KerPhi.T)
S[:ne,:ne] = np.diag(1.0/s[:ne])
kinv = np.dot(Vh.T, np.dot(S,U.T))
#kinv = np.linalg.pinv(cal.kpi.KerPhi)

bprms = [200.0, 90.0, 5.] # binary parameters


modl   = fit.binary_model(bprms, cal.kpi, cal.hdr) # binary model

# model-phase
# -----------
mphase = pc.phase_binary(cal.kpi.uv[:,0], cal.kpi.uv[:,1], 
                         cal.hdr['filter'], bprms)

# reconstructed phase
# -------------------
rphase = np.dot(kinv, modl) 

dxy = np.max(np.abs(cal.kpi.uv))
uvx = cal.kpi.uv[:,0]
uvy = cal.kpi.uv[:,1]
RED = cal.kpi.RED

#xs, ys = 100, 200
#xpi = np.linspace(0, -dxy, xs)
#ypi = np.linspace(-dxy, dxy, ys)

xs, ys = 200, 200
xpi = np.linspace(-dxy, dxy, xs)
ypi = np.linspace(-dxy, dxy, ys)

uvx = np.append(uvx, -uvx)
uvy = np.append(uvy, -uvy)
rphase = np.append(rphase, -rphase)
mphase = np.append(mphase, -mphase)
RED    = np.append(RED, RED)

reco = griddata((uvx, uvy), rphase, 
                (xpi[None,:], ypi[:,None]), method='linear')


orig = griddata((uvx, uvy), mphase, 
                (xpi[None,:], ypi[:,None]), method='linear')

ampl = griddata((uvx, uvy), RED, 
                (xpi[None,:], ypi[:,None]), method='linear')


plt.subplot(131)
plt.imshow(reco)
plt.title('Pseudo-inverse phase map')
plt.subplot(132)
plt.imshow(orig)
plt.title('Original phase map')
plt.subplot(133)
plt.imshow(ampl)
plt.title('Amplitude map')
